from datetime import datetime
from typing import TypedDict
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from langgraph.constants import START, END
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class WeatherAct(BaseModel):
    city: str = Field(description = "查询天气的城市名")
    date: str = Field(description = "查询天气的时间")

class WeatherState(TypedDict):
    query: str
    city: str
    date: str
    report: str
    answer: str

class WeatherAgent:

    def __init__(self, llm_model, tool):
        self.llm_model = llm_model
        self.tool = tool

    def extract_node(self, state: WeatherState):
        logger.info("WeatherAgent extract_node start")
        query = state["query"]
        extract_template = """
            你是一位问题分析高手，根据用户的输入问题提取其中的城市名和时间。
            用户输入的问题：{question}
        """
        extract_prompt = ChatPromptTemplate.from_template(extract_template)
        extract_chain = extract_prompt | self.llm_model.with_structured_output(WeatherAct)
        extract_result = extract_chain.invoke({ "question": query })
        logger.info(f"WeatherAgent extract_node extract_result: {extract_result}")
        return { "city": extract_result.city, "date": extract_result.date }

    def execute_node(self, state: WeatherState):
        logger.info("WeatherAgent execute_node start")
        city = state["city"]
        result = self.tool.func(city)
        logger.info(f"WeatherAgent execute_node result len: {len(result)}")
        return { "report": result }

    def generate_node(self, state: WeatherState):
        logger.info("WeatherAgent generate_node start")
        current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        generate_template = """
            当前时间：{current_date}
        
            你是一位问题分析高手，根据提供的上下文和用户的输入问题及提供的城市名和时间回答问题。
            
            上下文：{context}
            用户输入的问题：{question}
            城市名：{city}
            时间：{date}
            
            必须用中文详尽的回答完整的天气情况。
        """
        generate_prompt = ChatPromptTemplate.from_template(generate_template)
        generate_chain = generate_prompt | self.llm_model | StrOutputParser()
        generate_result = generate_chain.invoke({ "current_date": current_date, "context": state["report"], "question": state["query"], "city": state["city"], "date": state["date"] })
        logger.info(f"WeatherAgent generate_node generate_result len: {len(generate_result)}")
        return { "answer": generate_result }

    def build_graph(self):
        logger.info("WeatherAgent build_graph start")
        graph = StateGraph(WeatherState)
        graph.add_node("extract", self.extract_node)
        graph.add_node("execute", self.execute_node)
        graph.add_node("generate", self.generate_node)

        graph.add_edge(START, "extract")
        graph.add_edge("extract", "execute")
        graph.add_edge("execute", "generate")
        graph.add_edge("generate", END)
        return graph.compile()

    def invoke(self, query):
        logger.info(f"WeatherAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        logger.info(f"WeatherAgent invoke response: {response}")
        return { "document": response.get("report", None), "answer": response.get("answer", None) }
